﻿#define USE_TASKS

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;

#if USE_TASKS
using System.Threading;
using System.Threading.Tasks;
#endif

namespace TaskBuilder
{
    interface IProduce<T>
    {
        int Connect(Action<T> a);
        void Request(int id);
    }

    interface IConsumer
    {
        void Connect();
        void Request();
    }

    class IdMap<T>
    {
        private Dictionary<int, T> _im;
        private int _i;

        public IdMap()
        {
            _im = new Dictionary<int, T>();
        }

        public int Add(T otherId)
        {
#if USE_TASKS
            int id = Interlocked.Increment(ref _i);
#else
            int id = ++_i;
#endif

            _im.Add(id, otherId);

            return id;
        }

        public T this[int otherId]
        {
            get
            {
                return _im[otherId];
            }
        }

        public IEnumerable<T> Values
        {
            get { return _im.Values; }
        }
    }

    class Act<T> : IConsumer
    {
        private Action<T> _a;
        private IProduce<T> _p;
        private int _id;
        private int _n;

        public Act(Action<T> a, IProduce<T> p):
            this(a, 1, p)
        {
        }

        public Act(Action<T> a, int n, IProduce<T> p)
        {
            _a = a;
            _p = p;
            _n = n;
        }

        public void Connect()
        {
            _id = _p.Connect(InputReceived);
        }

        public void Request()
        {
            for (int n = 0; n < _n; n++)
                _p.Request(_id);
        }

        private void InputReceived(T t)
        {
#if USE_TASKS
            var task = new Task(Send, t);
            task.Start();
            task.ContinueWith(tk => Request());
#else
            _a(t);
            Request();
#endif
        }

        private void Send(object o)
        {
            _a((T)o);
        }
    }

    class Source<T> : IProduce<T>
    {
        private IEnumerator<T> _en;
        private Action<T> _a;
        private object _l = new object();

        public Source(IEnumerable<T> e)
        {
            _en = e.GetEnumerator();
        }

        public int Connect(Action<T> a)
        {
            _a = a;

            return 1;
        }

        public void Request(int id)
        {
#if USE_TASKS
            var t = new Task(Generate);
            t.Start();
#else
            Generate();
#endif
        }

        private void Generate()
        {
            T t;
            lock (_l)
            {
                if (!_en.MoveNext())
                    return;

                t = _en.Current;
            }
            _a(t);
        }
    }

    class Input<T> : IProduce<T>
    {
        private Action<T> _a;
        private object _l = new object();
        private int _n;
        private int _reqs;
        private Queue<T> _q = new Queue<T>();

        public Input()
            : this(int.MaxValue)
        {
        }
        public Input(int n)
        {
            _n = n;
        }

        public int Connect(Action<T> a)
        {
            _a = a;
            _reqs = 0;

            return 1;
        }

        public void Request(int id)
        {
            T t;
            lock (_l)
            {
                if (_q.Count == 0)
                {
                    _reqs++;
                    return;
                }
                t = _q.Dequeue();
            }
            _a(t);
        }

        public bool Write(T t)
        {
            lock (_l)
            {
                if (_reqs == 0)
                {
                    if (_q.Count == _n && _n != int.MaxValue)
                        return false;

                    _q.Enqueue(t);
                    return true;
                }
                _reqs--;
            }
            _a(t);
            return true;
        }
    }

    class Function<T, TResult> : IProduce<TResult>
    {
        private IProduce<T> _p;
        private Func<T, TResult> _f;
        private IdMap<int> _idm = new IdMap<int>();
        private Action<TResult> _a;

        public Function(Func<T, TResult> f, IProduce<T> p)
        {
            _p = p;
            _f = f;
        }

        public int Connect(Action<TResult> a)
        {
            if (_a != null)
                return 0;

            _a = a;

            return _idm.Add(_p.Connect(InputReceived));
        }

        public void Request(int id)
        {
            _p.Request(_idm[id]);
        }

        private void InputReceived(T t)
        {
#if USE_TASKS
            var task = new Task(Send, t);
            task.Start();
#else
            Send(t);
#endif
        }

        private void Send(object o)
        {
            TResult tr = _f((T)o);
            _a(tr);
        }
    }

    class JoinAll<T1, T2> : IProduce<Tuple<T1, T2>>
    {
        private IProduce<T1> _p1;
        private IProduce<T2> _p2;

        private Queue<T1> _q1;
        private Queue<T2> _q2;
        private object _l = new object();
        
        private IdMap<Tuple<int, int>> _idm = new IdMap<Tuple<int, int>>();
        private Action<Tuple<T1, T2>> _a;

        public JoinAll(IProduce<T1> p1, IProduce<T2> p2)
        {
            _p1 = p1;
            _p2 = p2;

            _q1 = new Queue<T1>();
            _q2 = new Queue<T2>();
        }

        public int Connect(Action<Tuple<T1, T2>> a)
        {
            if (_a != null)
                return 0;

            _a = a;

            return _idm.Add(new Tuple<int,int>(
                _p1.Connect(InputReceived1),
                _p2.Connect(InputReceived2)));
        }

        public void Request(int id)
        {
            Tuple<int, int> ids = _idm[id];
            _p1.Request(ids.Item1);
            _p2.Request(ids.Item2);
        }

        private void InputReceived1(T1 t)
        {
            Tuple<T1, T2> tp;
            lock (_l)
            {
                _q1.Enqueue(t);
                if (_q2.Count == 0)
                    return;
                tp = new Tuple<T1, T2>(_q1.Dequeue(), _q2.Dequeue());
            }
            _a(tp);
        }

        private void InputReceived2(T2 t)
        {
            Tuple<T1, T2> tp;
            lock (_l)
            {
                _q2.Enqueue(t);
                if (_q1.Count == 0)
                    return;
                tp = new Tuple<T1, T2>(_q1.Dequeue(), _q2.Dequeue());
            }
            _a(tp);
        }
    }

    class JoinAny<T1, T2> : IProduce<Tuple<T1, T2>>
    {
        private IProduce<T1> _p1;
        private IProduce<T2> _p2;

        private IdMap<Tuple<int, int>> _idm = new IdMap<Tuple<int, int>>();
        private Action<Tuple<T1, T2>> _a;

        public JoinAny(IProduce<T1> p1, IProduce<T2> p2)
        {
            _p1 = p1;
            _p2 = p2;
        }

        public int Connect(Action<Tuple<T1, T2>> a)
        {
            if (_a != null)
                return 0;

            _a = a;

            return _idm.Add(new Tuple<int, int>(
                _p1.Connect(InputReceived1),
                _p2.Connect(InputReceived2)));
        }

        public void Request(int id)
        {
            Tuple<int, int> ids = _idm[id];
            _p1.Request(ids.Item1);
            _p2.Request(ids.Item2);
        }

        private void InputReceived1(T1 t)
        {
            _a(new Tuple<T1, T2>(t, default(T2)));
        }

        private void InputReceived2(T2 t)
        {
            _a(new Tuple<T1, T2>(default(T1), t));
        }
    }

    class Branch<T> : IProduce<T>
    {
        private IProduce<T> _p;
        private int _id;
        private IdMap<Action<T>> _idm = new IdMap<Action<T>>();
        private int _req;
        private object _l = new object();

        public Branch(IProduce<T> p)
        {
            _p = p;
        }

        public int Connect(Action<T> a)
        {
            int id = _idm.Add(a);

            if (_id == 0)
                _id = _p.Connect(InputReceived);

            return id;
        }

        public void Request(int id)
        {
            lock (_l)
            {
                _req++;

                if (_req <= 0)
                    return;

                _req -= _idm.Values.Count();
            }
            _p.Request(_id);
        }

        private void InputReceived(T t)
        {
            foreach (Action<T> a in _idm.Values)
            {
#if USE_TASKS
                var task = new Task(Send, (object)(new object[] { t, a }));
                task.Start();
#else
                a(t);
#endif
            }
        }

        private void Send(object o)
        {
            object[] oa = (object[])o;
            ((Action<T>)oa[1])((T)oa[0]);
        }
    }

    class Collect<T> : IProduce<T[]>
    {
        private IProduce<T> _p;
        private int _n;
        private IdMap<int> _idm = new IdMap<int>();
        private Queue<T> _q = new Queue<T>();
        private object _l = new object();
        private Action<T[]> _a;

        public Collect(int n, IProduce<T> p)
        {
            _p = p;
            _n = n;
        }

        public int Connect(Action<T[]> a)
        {
            if (_a != null)
                return 0;

            _a = a;

            return _idm.Add(_p.Connect(InputReceived));
        }

        public void Request(int id)
        {
            _p.Request(_idm[id]);
        }

        private void InputReceived(T t)
        {
            T[] ta;
            lock (_l)
            {
                _q.Enqueue(t);
                if (_q.Count >= _n)
                {
                    ta = new T[_n];
                    for (int i = 0; i < ta.Length; i++)
                        ta[i] = _q.Dequeue();
                }
                else
                {
                    ta = null;
                }
            }

            if (ta != null)
                _a(ta);
            else
                _p.Request(0);
        }
    }

    class Separate<T> : IProduce<T>
    {
        private IProduce<T[]> _p;
        private Queue<T> _q;
        private object _l = new object();
        private Action<T> _a;
        private IdMap<int> _idm = new IdMap<int>();

        public Separate(IProduce<T[]> p)
        {
            _p = p;
            _q = new Queue<T>();
        }

        public int Connect(Action<T> a)
        {
            _a = a;

            return _idm.Add(_p.Connect(InputReceived));
        }

        public void Request(int id)
        {
            bool req;
            T t = default(T);
            lock (_l)
            {
                req = _q.Count == 0;
                if (!req)
                    t = _q.Dequeue();
            }

            if (req)
                _p.Request(_idm[id]);
            else
                _a(t);
        }

        private void InputReceived(T[] ta)
        {
            T t;
            lock (_l)
            {
                foreach (T ct in ta)
                    _q.Enqueue(ct);

                t = _q.Dequeue();
            }
            _a(t);
        }
    }

    class GroupAll<T> : IProduce<T[]>
    {
        private IProduce<T>[] _ps;
        private Queue<T>[] _qs;
        private IdMap<int[]> _idm = new IdMap<int[]>();
        private Action<T[]> _a;
        private object _l = new object();

        public GroupAll(params IProduce<T>[] ps)
        {
            _ps = ps;

            _qs = new Queue<T>[ps.Length];
            for (int i = 0; i < ps.Length; i++)
                _qs[i] = new Queue<T>();
        }

        public int Connect(Action<T[]> a)
        {
            if (_a != null)
                return 0;

            _a = a;

            int[] ids = new int[_ps.Length];

            for (int i = 0; i < _ps.Length; i++)
            {
                int qi = i;
                ids[i] = _ps[i].Connect(t => InputReceived(qi, t));
            }

            return _idm.Add(ids);
        }

        public void Request(int id)
        {
            for (int i = 0; i < _ps.Length; i++)
                _ps[i].Request(_idm[id][i]);
        }

        private void InputReceived(int n, T t)
        {
            T[] ta = new T[_qs.Length];
            lock (_l)
            {
                _qs[n].Enqueue(t);
                foreach (Queue<T> q in _qs)
                    if (q.Count == 0)
                        return;

                for (int j = 0; j < _qs.Length; j++)
                    ta[j] = _qs[j].Dequeue();
            }
            _a(ta);
        }
    }

    class GroupAny<T> : IProduce<T>
    {
        private IProduce<T>[] _ps;
        private IdMap<int[]> _idm = new IdMap<int[]>();
        private Action<T> _a;

        public GroupAny(params IProduce<T>[] ps)
        {
            _ps = ps;
        }

        public int Connect(Action<T> a)
        {
            if (_a != null)
                return 0;

            _a = a;

            int[] ids = new int[_ps.Length];

            for (int i = 0; i < _ps.Length; i++)
            {
                int qi = i;
                ids[i] = _ps[i].Connect(t => InputReceived(qi, t));
            }

            return _idm.Add(ids);
        }

        public void Request(int id)
        {
            for (int i = 0; i < _ps.Length; i++)
                _ps[i].Request(_idm[id][i]);
        }

        private void InputReceived(int n, T t)
        {
            _a(t);
        }
    }

    class Ungroup<T> : IProduce<T>
    {
        private int _pi;
        private int _n;
        private IProduce<T[]> _p;
        private IdMap<int> _idm = new IdMap<int>();
        private Action<T> _a;

        public Ungroup(int n, int pi, IProduce<T[]> p)
        {
            _n = n;
            _pi = pi;
            _p = p;
        }

        public int Connect(Action<T> a)
        {
            if (_a != null)
                return 0;

            _a = a;

            return _idm.Add(_p.Connect(InputReceived));
        }

        public void Request(int id)
        {
            _p.Request(_idm[id]);
        }

        private void InputReceived(T[] ta)
        {
            Debug.Assert(ta.Length % _n == 0,
                string.Format("Cannot split input size {0} across {1} parallel instances.", ta.Length, _n));
            for (int i = _pi; i < ta.Length; i += _n)
                _a(ta[i]);
        }
    }

    class Point<T> : IProduce<T>
    {
        private IProduce<T> _p;

        public void Attach(IProduce<T> p)
        {
            _p = p;
        }

        public int Connect(Action<T> a)
        {
            return _p.Connect(a);
        }

        public void Request(int id)
        {
            _p.Request(id);
        }
    }

    static class TB
    {
        private class EvalScope : IDisposable
        {
            private List<IConsumer> _cs = new List<IConsumer>();

            public void Add(IConsumer c)
            {
                _cs.Add(c);
            }

            public void Dispose()
            {
                foreach (IConsumer c in _cs)
                    c.Connect();

                foreach (IConsumer c in _cs)
                {
#if USE_TASKS
                    var t = new Task(c.Request);
                    t.Start();
#else
                    c.Request();
#endif
                }

                _currentScope = null;
            }
        }

        [ThreadStatic]
        private static EvalScope _currentScope;

        public static IDisposable Scope()
        {
            if (_currentScope != null)
                throw new InvalidOperationException("Nested scopes are not supported.");

            _currentScope = new EvalScope();
            return _currentScope;
        }

        public static void Act<T>(Action<T> a, IProduce<T> p)
        {
            if (_currentScope == null)
                throw new InvalidOperationException("Cannot use PG.Act outside an evaluation scope.");

            _currentScope.Add(new Act<T>(a, p));
        }

        public static void Act<T>(Action<T> a, int n, IProduce<T> p)
        {
            if (_currentScope == null)
                throw new InvalidOperationException("Cannot use PG.Act outside an evaluation scope.");

            _currentScope.Add(new Act<T>(a, n, p));
        }

        public static void Dispatch<T>(Action<T> a, IProduce<T> p)
        {
            if (_currentScope == null)
                throw new InvalidOperationException("Cannot use PG.Act outside an evaluation scope.");

            _currentScope.Add(new Act<T>(a, p));
        }

        public static IProduce<T> Func<T>(IEnumerable<T> e)
        {
            return new Source<T>(e);
        }

        public static IProduce<T> Func<T>(Func<IEnumerable<T>> e)
        {
            return new Source<T>(e());
        }

        public static IProduce<TResult> Func<T, TResult>(Func<T, TResult> f, IProduce<T> p)
        {
            return new Function<T, TResult>(f, p);
        }

        public static IProduce<TResult[]> Func<T, TResult>(Func<T, TResult> f, int n, IProduce<T[]> p)
        {
            IProduce<T>[] t = TB.Ungroup(n, p);
            IProduce<TResult>[] tr = new IProduce<TResult>[n];
            for (int i = 0; i < n; i++)
                tr[i] = new Function<T, TResult>(f, t[i]);
            return TB.GroupAll(tr);
        }

        public static IProduce<Tuple<T1, T2>> All<T1, T2>(IProduce<T1> p1, IProduce<T2> p2)
        {
            return new JoinAll<T1, T2>(p1, p2);
        }

        public static IProduce<Tuple<T1, T2>> Any<T1, T2>(IProduce<T1> p1, IProduce<T2> p2)
        {
            return new JoinAny<T1, T2>(p1, p2);
        }

        public static IProduce<T> Branch<T>(IProduce<T> p)
        {
            return new Branch<T>(p);
        }

        // Consume Ts over time from one producer and produce a T[] from one producer
        public static IProduce<T[]> Collect<T>(int n, IProduce<T> p)
        {
            return new Collect<T>(n, p);
        }

        // Consume a T[] from one producer and produce Ts over time from one producer
        public static IProduce<T> Separate<T>(IProduce<T[]> p)
        {
            return new Separate<T>(p);
        }

        // Group all Ts from several producers and produce a T[]
        public static IProduce<T[]> GroupAll<T>(params IProduce<T>[] ps)
        {
            return new GroupAll<T>(ps);
        }

        // Group any T from several producers and produce a T[]
        public static IProduce<T> GroupAny<T>(params IProduce<T>[] ps)
        {
            return new GroupAny<T>(ps);
        }

        // Consume a T[] and produce Ts from several producers
        public static IProduce<T>[] Ungroup<T>(int n, IProduce<T[]> p)
        {
            var b = new Branch<T[]>(p);
            var pa = new IProduce<T>[n];
            for (int i = 0; i < n; i++)
                pa[i] = new Ungroup<T>(n, i, b);

            return pa;
        }

        public static Point<T> Point<T>()
        {
            return new Point<T>();
        }

        public static Input<T> Input<T>()
        {
            return new Input<T>();
        }
    }
}
